"""
author: liuxu
date: 2024.12.3
description: pmc硬件计数器功能测试
"""
import os.path
import re
import allure
import pytest
import csv
import glob
import random


from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd
from base_public.testbase import BasePublic
from conftest import dcu_info


# pmc指标简称
basic_short_names = ['kernel-name', 'dispatch', 'gpu-id', 'queue-id', 'queue-index', 'pid', 'tid', 'grd', 'wgr',
                              'lds', 'scr', 'vgpr', 'sgpr', 'fbar', 'sig', 'TIME', 'VALUBusy', 'performance','LDSInsts',
                              'LDSBankConflict', 'L1MemoryBusy', 'L1MemoryStalled', 'L2WriteUnitStalled', 'L2CacheHit']
read_short_names = ['kernel-name', 'dispatch', 'gpu-id', 'queue-id', 'queue-index', 'pid', 'tid', 'grd', 'wgr',
                              'lds', 'scr', 'vgpr', 'sgpr', 'fbar', 'sig', 'TIME', 'VALUBusy', 'performance','LDSInsts',
                              'LDSBankConflict', 'L1MemoryBusy', 'L1MemoryStalled', 'L2ReadReqs', 'SizeOfL2Read']
write_short_names = ['kernel-name', 'dispatch', 'gpu-id', 'queue-id', 'queue-index', 'pid', 'tid', 'grd', 'wgr',
                              'lds', 'scr', 'vgpr', 'sgpr', 'fbar', 'sig', 'TIME', 'VALUBusy', 'performance','LDSInsts',
                              'LDSBankConflict', 'L1MemoryBusy', 'L1MemoryStalled', 'L2WriteReqs', 'SizeOfL2Write']
# pmc原始数据指标（k为指标名称，v为个数）
PMC_DATA_DICT = {'GRBM_COUNT': 1, 'GRBM_GUI_ACTIVE': 1, 'SQ_ACTIVE_INST_VALU': 1, 'SQ_INSTS_FLAT_LDS_ONLY': 1,
                 'SQ_INSTS_LDS': 1, 'SQ_INSTS_VALU': 1, 'SQ_INSTS_VMEM_RD': 1, 'SQ_INSTS_VMEM_WR': 1,
                 'SQ_LDS_BANK_CONFLICT': 1, 'SQ_WAIT_INST_LDS': 1, 'TA_FLAT_READ_WAVEFRONTS': 15, 'TA_TA_BUSY': 15,
                 'TCC_EA_WRREQ_STALL': 15, 'TCC_EA1_WRREQ_STALL': 15, 'TCC_HIT': 15, 'TCC_MISS': 15,
                 'TCP_TCP_TA_DATA_STALL_CYCLES': 15}

@pytest.mark.high
@allure.feature('PMC功能测试')
class TestPMC(BasePublic):
    def _get_pmc_file_name(self, output, file_type="txt"):
        """
        获取输出的pmc文件名称
        :param output:
        :param file_type: txt/csv
        :return:
        """
        prof_pid = self.get_prof_pid(output)
        if file_type == "txt":
            match_name = f"pmc_results_{prof_pid}.txt"
        else:
            match_name = f"pmc_results_{prof_pid}*.csv"  # csv文件可能有一个或多个
        pmc_file_name_l = glob.glob(match_name)
        assert len(pmc_file_name_l) >= 1, "The pmc data file was not generated"
        return pmc_file_name_l[0]

    def _get_unopt_matrix_exe(self):
        # pmc使用-O0和-O3的结果会有差异,优化后性能会变好
        return self.compile_source_code(src_file_name="hipSimple/matrix.cpp",
                                            des_exe_name="matrix-O0-g",
                                            compile_paras="-O0 -g")
    @pytest.mark.smoke
    def test_basic_pmc(self, test_env, dcu_info):
        """
        获取pmc基础的指标（默认）
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据
        out_pmc_file = self._get_pmc_file_name(output)
        pmc_info = self.get_pmc_info_from_txt(out_pmc_file, pmc_type="basic")

        # 校验pmc指标信息
        assert len(pmc_info) == 1
        kernel_name, indecator_d = pmc_info[0].popitem()
        assert "matrixTranspose" in kernel_name
        assert indecator_d.get("kernel_dispatch_index") == 0
        assert indecator_d.get("gpu_id") == dcu_info.get("devices_id")[0]
        assert indecator_d.get("dispatch_queue_id") == 0
        assert indecator_d.get("dispatch_queue_index") >= 0
        # MAX_PID = 65536  # cat /proc/sys/kernel/pid_max
        MAX_PID = self.get_sys_max_pid()
        assert  self.get_prof_pid(output) <= indecator_d.get("process_id") == indecator_d.get("thread_id") < MAX_PID
        assert indecator_d.get("grid_size") == 1048576
        assert indecator_d.get("work_group_size") == 16
        assert indecator_d.get("shared_memory_size") == 0
        assert indecator_d.get("scratch_size") == 104
        assert indecator_d.get("vgpr_count") in (12, 24)
        assert indecator_d.get("sgpr_count") == 48
        # assert indecator_d.get("fbarrier_count") == 624587557  # dtk-25.10删除了fbarrier_count指标
        assert indecator_d.get("signal_handle") == 0
        assert indecator_d.get("kernel_time") < 0.001  # 小于1ms
        assert 1.0 < indecator_d.get("processed_ALU_instructions") < 15.0
        assert 300.0 < indecator_d.get("performance") < 1400.0
        assert indecator_d.get("shared_memory_operation") == 1048576
        assert indecator_d.get("shared_memory_bank_conflict") == 0.0
        assert 20.0 < indecator_d.get("L1_cache_unit_is_active") < 40.0
        assert 5.0 < indecator_d.get("L1_cache_unit_is_stalled") < 20.0
        assert 0.05 < indecator_d.get("L2_cache_write_unit_is_stalled") < 0.9
        assert 80.0 < indecator_d.get("L2_cache_hit_rate") < 90.0

    def test_read_pmc(self, test_env):
        """
        获取pmc读相关的指标
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-read {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据
        out_pmc_file = self._get_pmc_file_name(output)
        pmc_info = self.get_pmc_info_from_txt(out_pmc_file, pmc_type="read")

        # 校验获取的pmc信息
        assert len(pmc_info) == 1
        kernel_name, indecator_d = pmc_info[0].popitem()
        assert 65100.0 < indecator_d.get("number_of_L2_cache_read_requests") < 66900.0
        assert 4050.0 < indecator_d.get("size_of_L2_cache_read") < 4190.0

    def test_write_pmc(self, test_env):
        """
        获取pmc写相关的指标
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-write {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据
        out_pmc_file = self._get_pmc_file_name(output)
        pmc_info = self.get_pmc_info_from_txt(out_pmc_file, pmc_type="write")

        # 校验获取的pmc信息
        assert len(pmc_info) == 1
        kernel_name, indecator_d = pmc_info[0].popitem()
        assert 90000.0 < indecator_d.get("number_of_L2_cache_write_requests") < 234000.0
        assert 5000.0 < indecator_d.get("size_of_L2_cache_write") < 15000.0

    def test_pmc_fullname_indicator(self, test_env):
        """
        获取pmc指标名称类型为全称(默认)
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-type 0 {exe_path}", cwd=test_env)
        assert ret is True

    @pytest.mark.parametrize("pmc_para, expected_indicator_l", (
    ['--pmc', basic_short_names], ["--pmc-read", read_short_names], ["--pmc-write", write_short_names]))
    def test_pmc_abbreviation_indicator(self, test_env, pmc_para, expected_indicator_l):
        """
        获取pmc指标名称类型为简称
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc {pmc_para} --pmc-type 1 {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据,校验指标简称值
        pmc_file_info_l = self.read_text_file(file_name=self._get_pmc_file_name(output)).rstrip("\n").split("\n")  # 去除文件末尾空行
        for index, expected_indicator in enumerate(expected_indicator_l):
            assert expected_indicator in pmc_file_info_l[index]  # 按顺序校验指标名称

    def test_pmc_original_indicator(self, test_env, dcu_info):
        """
        获取pmc原始数据（pmc数据为原始数据）
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name='multiproc_case/multithreading_singledevice.cpp') # 3线程单卡

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-type 2 {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc原始数据并进行校验
        pmc_data = self.get_pmc_original_data_from_txt(self._get_pmc_file_name(output))
        assert len(pmc_data) == 3  # 核函数个数
        for idx, kernel_d in enumerate(pmc_data):
            # （1）校验kernel dispatch信息
            assert kernel_d.get("dispatch_index") == idx
            assert "computeKernel" in kernel_d.get("kernel_name")
            assert kernel_d.get("gpu-id") == dcu_info["devices_id"][-1]

            # （2）校验计数器信息
            # assert kernel_d["GRBM_COUNT"] == kernel_d["GRBM_GUI_ACTIVE"] > 0
            assert kernel_d["GRBM_COUNT"] > 0 and kernel_d["GRBM_GUI_ACTIVE"] > 0
            for indicator_name, indicator_num in PMC_DATA_DICT.items():  # 校验指标都存在
                if indicator_num == 1:
                    assert kernel_d.get(indicator_name) is not None, f"The indicator '{indicator_name}' was not found."
                else:
                    for i in range(indicator_num):
                        assert i in kernel_d.get(indicator_name), f"The indicator {indicator_name}{indicator_num} was not found."

            # （3）校验时间戳
            assert 0 < kernel_d["DispatchNs"] < kernel_d["BeginNs"] < kernel_d["EndNs"] < kernel_d["CompleteNs"]

    def test_pmc_csv_indicator(self, test_env, dcu_info):
        """
        获取pmc指标名称类型为csv表格
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name='multiproc_case/multithreading_singledevice.cpp') # 3线程单卡

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-type 3 {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据并进行校验
        pmc_data = self.get_pmc_info_from_csv(self._get_pmc_file_name(output, file_type="csv"))
        assert len(pmc_data) == 3  # 核函数个数
        for idx, kernel_d in enumerate(pmc_data):
            # （1）校验kernel dispatch信息
            assert int(kernel_d.get("Index")) == idx
            assert "computeKernel" in kernel_d.get("KernelName")
            assert int(kernel_d.get("gpu-id")) == dcu_info["devices_id"][-1]

            # （2）校验计数器信息
            assert kernel_d["GRBM_COUNT"], kernel_d["GRBM_GUI_ACTIVE"] > 0
            for indicator_name, indicator_num in PMC_DATA_DICT.items():  # 校验PMC原始数据指标都存在
                if indicator_num == 1:
                    assert kernel_d.get(indicator_name) is not None, f"The indicator '{indicator_name}' was not found."
                else:
                    for i in range(indicator_num):
                        assert kernel_d.get(f"{indicator_name}[{indicator_num}]") >= 0, \
                            f"The indicator {indicator_name}{indicator_num} was not found."

            # （3）校验时间戳
            assert 0 < kernel_d["DispatchNs"] < kernel_d["BeginNs"] < kernel_d["EndNs"] < kernel_d["CompleteNs"]
            assert kernel_d["CompleteNs"] - kernel_d["DispatchNs"] < 2 * 10**9  # 校验下发到结束小于1s

    def test_pmc_filter_kernelname(self, test_env):
        """
        获取pmc指标,过滤kernel名称
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="vcopy_test/vcopy.cpp")

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --kernel-name vecCopy_2 {exe_path} -n 1048576 -b 256 -i 1 --multikernel", cwd=test_env)
        assert ret is True

        # 获取pmc文件内容，只包括过滤的pmc信息
        pmc_content = self.read_text_file(self._get_pmc_file_name(output))
        assert "vecCopy_2" in pmc_content
        assert "vecCopy_1" not in pmc_content

    @pytest.mark.exception
    def test_pmc_filter_invalid_kernelname(self, test_env):
        """
        测试过滤无效的kernel名称
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="vcopy_test/vcopy.cpp")

        # 执行测试命令
        ret, output = exec_cmd(
            f"hipprof --pmc --kernel-name vecCopy_3 {exe_path} -n 1048576 -b 256 -i 1 --multikernel", cwd=test_env)
        assert ret is True

        # 校验未生成pmc文件
        assert self.exist_file(self._get_pmc_file_name(output))

    def test_pmc_list_basic (self, test_env):
        """
        获取pmc支持的硬件计数器指标
        :return:
        """
        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --list-basic", cwd=test_env)

        # 校验结果
        assert ret is True
        assert "Basic HW counters:" in output

    def test_pmc_list_derived (self, test_env):
        """
        获取pmc所有派生指标及公式
        :return:
        """
        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --list-derived", cwd=test_env)

        # 校验结果
        assert ret is True
        assert "Derived metrics:" in output

    @pytest.mark.smoke
    def test_pmc_filter_indicator(self, test_env, dcu_info):
        """
        测试pmc过滤指标(当使用-i参数时，默认生成csv文件)
        :return:
        """
        # 随机测试2张可用卡
        avail_device_ids = [str(i) for i in dcu_info.get("devices_id")]
        test_device_ids =  random.sample(avail_device_ids, 2)
        mylogger.debug(f"test_device_ids: {test_device_ids}")

        # 判断当前环境是否为bw
        IS_BW = dcu_info.get('gfx_arch') == "gfx936"

        # 过滤的pmc指标
        if IS_BW:
            # bw上TCC_REQ_sum > TCC_HIT_sum + TCC_MISS_sum
            PMC_Ind = "TCC_REQ_sum,TCC_READ_sum,TCC_WRITE_sum"
        else:
            PMC_Ind = "GRBM_COUNT,SQ_CYCLES,GRBM_GUI_ACTIVE,SQ_INSTS_FLAT,SQ_INSTS_SALU,SQ_LDS_BANK_CONFLICT,SQ_WAVES,TCC_REQ_sum,TCC_HIT_sum,TCC_MISS_sum"
        cfg_name = os.path.join(test_env, "filter_config.txt")
        filter_info = f"""pmc: {PMC_Ind}
        range: 0:3
        gpu: {' '.join(test_device_ids)} 
        kernel: vectorAdd simple"""

        # 编译待测用例
        case_path = self.compile_source_code(src_file_name="multiproc_case/multiprocess_multithread_multidevice_test.cpp",
                                             des_exe_name="multiproc_case/multiprocess_multithread_multidevice_test-O0-g",
                                             compile_paras="-O0 -g")
        assert self.create_text_file(path=cfg_name, content=filter_info) is True

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc -o ./tmp/test-redirect_name -i {cfg_name} {case_path}", cwd=test_env)
        assert ret is True, "Failed to execute the pmc command"
        assert "HIP_PROF:finish" in output.strip().split("\n")[-1]
        assert re.search("fail|error|can not", output, re.DOTALL) is None
        assert self.exist_file("./tmp/test-redirect_name.csv") is True, "not found pmc output file"

        # 获取pmc数据并进行校验
        pmc_data_l = self.get_pmc_info_from_csv(os.path.join(test_env, "./tmp/test-redirect_name.csv"))
        for kernel_pmc_d in pmc_data_l:
            # 校验csv中包含所有过滤的pmc指标
            assert set(PMC_Ind.split(",")).issubset(kernel_pmc_d.keys())
            # 校验指标大小值
            if IS_BW:
                # TCC_REQ_sum = TCC_READ_sum + TCC_WRITE_sum
                assert int(kernel_pmc_d.get("TCC_REQ_sum")) == int(kernel_pmc_d.get("TCC_READ_sum")) + int(
                    kernel_pmc_d.get("TCC_WRITE_sum"))
            else:
                # TCC_REQ_sum = TCC_HIT_sum + TCC_MISS_sum
                assert int(kernel_pmc_d.get("TCC_REQ_sum")) == int(kernel_pmc_d.get("TCC_HIT_sum")) + int(
                    kernel_pmc_d.get("TCC_MISS_sum"))

    @pytest.mark.high
    @pytest.mark.parametrize("input_para, comp_opt_level", (
            ['--pmc', "-O0"], ["--pmc-read", "-O3"], ["--pmc-write", "-O3"]))
    def test_pmc_kernel_opt(self, test_env, dcu_info, input_para, comp_opt_level):
        """
        测试kernel优化前后的pmc指标(每一项指标都不为0)
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name='pmc_sqtt/pmc_optimization_test.cpp',  # kernel优化前后对比
                                            des_exe_name="pmc_optimization_test-O0-g",
                                            compile_paras=f"{comp_opt_level} -g")
        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc {input_para} {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc数据
        out_pmc_file = self._get_pmc_file_name(output)
        pmc_type = {"--pmc-read": "read", "--pmc-write": "write"}.get(input_para, "basic")
        pmc_info = self.get_pmc_info_from_txt(out_pmc_file, pmc_type=pmc_type)
        mylogger.info(pmc_info)

        # 校验对比2个kernel的pmc指标信息
        assert len(pmc_info) == 2
        base_kernel_d, opt_kernel_d = pmc_info
        assert "baseline_kernel" in next(iter(base_kernel_d.keys())) and "optimized_kernel" in next(iter(opt_kernel_d.keys()))

        base_kernel_indicator_d = next(iter(base_kernel_d.values()))
        opt_kernel_indicator_d = next(iter(opt_kernel_d.values()))
        assert base_kernel_indicator_d["grid_size"] == opt_kernel_indicator_d["grid_size"] == 262144
        assert base_kernel_indicator_d["work_group_size"] == opt_kernel_indicator_d["work_group_size"] == 256
        assert 0 < base_kernel_indicator_d["shared_memory_size"], opt_kernel_indicator_d["shared_memory_size"]
        assert 0 < base_kernel_indicator_d["scratch_size"], opt_kernel_indicator_d["scratch_size"]
        # assert 0 < base_kernel_indicator_d["fbarrier_count"], opt_kernel_indicator_d["fbarrier_count"]  # dtk-25.10删除了fbarrier_count指标
        assert 0 < base_kernel_indicator_d["shared_memory_operation"], opt_kernel_indicator_d["shared_memory_operation"]
        assert base_kernel_indicator_d["shared_memory_bank_conflict"] > opt_kernel_indicator_d["shared_memory_bank_conflict"] * 10 >= 0
        # 优化后kernel执行时间大幅缩减
        assert base_kernel_indicator_d["kernel_time"] > opt_kernel_indicator_d["kernel_time"] * 3 > 0

    @pytest.mark.medium
    def test_pmc_off_always(self, test_env):
        """
        全程关闭pmc抓取（未通过api开启）
        :return:
        """
        # 编译待测用例
        exe_path = self._get_unopt_matrix_exe()

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --pmc --pmc-off {exe_path}", cwd=test_env)
        assert ret is True

        # 校验未生成pmc文件，或者生成文件为空
        pmc_file_path = self._get_pmc_file_name(output)
        assert not self.exist_file(pmc_file_path) or len(self.read_text_file(pmc_file_path).strip()) == 0

    def test_pmc_off(self, test_env):
        """
        关闭pmc抓取（通过hipProfilerStart/Stop接口控制分析范围）
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name='pmc_sqtt/pmc_optimization_test.cpp') # 通过hipProfilerStart/Stop控制只分析优化后kernel

        # 执行测试命令
        ret, output = exec_cmd(f"hipprof --hip-trace --pmc --pmc-off {exe_path}", cwd=test_env)
        assert ret is True

        # 获取pmc信息
        pmc_file_path = self._get_pmc_file_name(output)
        pmc_info = self.get_pmc_info_from_txt(pmc_file_path)

        # 校验只生成hipProfilerStart/Stop范围内核函数的pmc信息
        assert len(pmc_info) == 1
        kernel = next(iter(pmc_info[0]))
        assert "optimized_kernel" in kernel

    @pytest.mark.medium
    def test_pmc_empty_kernel(self, test_env):  # http://hpczentao.sugon.com/bug-view-106127.html
        """
        测试分析空kernel的pmc数据
        :param test_env:
        :return:
        """
        # 编译待测用例
        exe_path = self.compile_source_code(src_file_name="rtc/rtc_perf_test.cpp")

        # 执行测试命令
        expected_execution_time = 5  # 秒
        ret, output = exec_cmd(f"hipprof --pmc --hip-trace --trace-args {exe_path} 0 1", cwd=test_env,
                               timeout=expected_execution_time)
        assert ret is True

        # 校验api trace数据
        assert self.check_trace_output_success(output)

        # 获取pmc信息
        pmc_file_path = self._get_pmc_file_name(output)
        pmc_info_l = self.get_pmc_info_from_txt(pmc_file_path)

        # 校验pmc数据是否正常
        assert len(pmc_info_l) == 1 and "emptyKernel" in pmc_info_l[0]
        kernel_info_d = pmc_info_l[0]
        assert int(kernel_info_d.get("emptyKernel").get("kernel_time")) < 0.000005  # 空核函数耗时不超过5us






        
